//
//  CTCollectionViewColumnLayout.m
//
//  Created by Andney Jack on 2021/6/2.
//

#import "CTCollectionViewColumnLayout.h"
#import "CTCollectionViewSectionInfo.h"

NSString *const CTCollectionElementKindSectionHeader = @"CTCollectionViewColumnLayoutSectionHeader";
NSString *const CTCollectionElementKindSectionFooter = @"CTCollectionViewColumnLayoutSectionFooter";

@interface CTCollectionViewColumnLayout ()
{
    BOOL _haveFloatingHeaderOrFooter;
}
@property (nonatomic, weak) id <CTCollectionViewDelegateColumnLayout> delegate;
@property (nonatomic, strong) NSMutableArray<CTCollectionViewSectionInfo *> *sectionInfos;

@end

@implementation CTCollectionViewColumnLayout

static CGFloat CTFloorCGFloat(CGFloat value)
{
    CGFloat scale = [UIScreen mainScreen].scale;
    return floor(value * scale) / scale;
}

#pragma mark - Public Accessors

- (void)setColumnCount:(NSInteger)columnCount {
    if (_columnCount != columnCount) {
        _columnCount = columnCount;
        [self invalidateLayout];
    }
}

- (void)setColumnSpacing:(CGFloat)columnSpacing {
    if (_columnSpacing != columnSpacing) {
        _columnSpacing = columnSpacing;
        [self invalidateLayout];
    }
}

- (void)setInteritemSpacing:(CGFloat)interitemSpacing {
    if (_interitemSpacing != interitemSpacing) {
        _interitemSpacing = interitemSpacing;
        [self invalidateLayout];
    }
}

- (void)setHeaderHeight:(CGFloat)headerHeight {
    if (_headerHeight != headerHeight) {
        _headerHeight = headerHeight;
        [self invalidateLayout];
    }
}

- (void)setFooterHeight:(CGFloat)footerHeight {
    if (_footerHeight != footerHeight) {
        _footerHeight = footerHeight;
        [self invalidateLayout];
    }
}

- (void)setHeaderInset:(UIEdgeInsets)headerInset {
    if (!UIEdgeInsetsEqualToEdgeInsets(_headerInset, headerInset)) {
        _headerInset = headerInset;
        [self invalidateLayout];
    }
}

- (void)setFooterInset:(UIEdgeInsets)footerInset {
    if (!UIEdgeInsetsEqualToEdgeInsets(_footerInset, footerInset)) {
        _footerInset = footerInset;
        [self invalidateLayout];
    }
}

- (void)setSectionInset:(UIEdgeInsets)sectionInset {
    if (!UIEdgeInsetsEqualToEdgeInsets(_sectionInset, sectionInset)) {
        _sectionInset = sectionInset;
        [self invalidateLayout];
    }
}

- (void)setItemRenderDirection:(CTCollectionViewItemRenderDirection)itemRenderDirection {
    if (_itemRenderDirection != itemRenderDirection) {
        _itemRenderDirection = itemRenderDirection;
        [self invalidateLayout];
    }
}

- (void)setAllowUseSafeAreaInsets:(BOOL)autoChangeSafeAreaInset {
    if (_allowUseSafeAreaInsets != autoChangeSafeAreaInset) {
        _allowUseSafeAreaInsets = autoChangeSafeAreaInset;
        [self invalidateLayout];
    }
}

- (void)setSectionHeadersPinToVisibleBounds:(BOOL)sectionHeadersPinToVisibleBounds {
    if (_sectionHeadersPinToVisibleBounds != sectionHeadersPinToVisibleBounds) {
        _sectionHeadersPinToVisibleBounds = sectionHeadersPinToVisibleBounds;
        [self invalidateLayout];
    }
}

- (void)setSectionFootersPinToVisibleBounds:(BOOL)sectionFootersPinToVisibleBounds {
    if (_sectionFootersPinToVisibleBounds != sectionFootersPinToVisibleBounds) {
        _sectionFootersPinToVisibleBounds = sectionFootersPinToVisibleBounds;
        [self invalidateLayout];
    }
}

- (NSInteger)columnCountForSection:(NSInteger)section {
    return [self m_columnCountForSection:section];
}

- (CGFloat)contentWidthForSection:(NSInteger)section {
    UIEdgeInsets sectionInset = [self m_insetForSectionAtIndex:section];
    CGFloat contentWidth = self.collectionView.bounds.size.width - sectionInset.left - sectionInset.right;
    if (self.allowUseSafeAreaInsets) {
        if (@available(iOS 11.0, *)) {
            UIEdgeInsets safeAreaInsets = self.collectionView.safeAreaInsets;
            contentWidth -= safeAreaInsets.left + safeAreaInsets.right;
        }
    }
    return MAX(contentWidth, 0.0);
}

- (CGFloat)itemWidthForSection:(NSInteger)section {
    CGFloat contentWidth = [self contentWidthForSection:section];
    NSInteger columnCount = [self m_columnCountForSection:section];
    CGFloat columnSpacing = [self m_columnSpacingForSectionAtIndex:section];

    return CTFloorCGFloat((contentWidth - (columnCount - 1) * columnSpacing) / columnCount);
}

#pragma mark - Private Accessors

- (NSMutableArray<CTCollectionViewSectionInfo *> *)sectionInfos {
    if (!_sectionInfos) {
        _sectionInfos = [[NSMutableArray alloc] initWithCapacity:0];
    }
    return _sectionInfos;
}

- (id <CTCollectionViewDelegateColumnLayout>)delegate {
    return (id <CTCollectionViewDelegateColumnLayout>)self.collectionView.delegate;
}

- (NSInteger)m_columnCountForSection:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:columnCountForSection:)]) {
        return [self.delegate collectionView:self.collectionView layout:self columnCountForSection:section];
    } else {
        return self.columnCount;
    }
}

- (CGFloat)m_heightForItemAtIndexPath:(NSIndexPath *)indexPath columnWidth:(CGFloat)columnWidth
{
    if (!self.collectionView) {
        return 0.0f;
    }

    if ([self.delegate respondsToSelector:@selector(collectionView:layout:sizeForItemAtIndexPath:columnWidth:)]) {
        CGSize itemSize = [self.delegate collectionView:self.collectionView layout:self sizeForItemAtIndexPath:indexPath columnWidth:columnWidth];
        if (itemSize.width > 0 && itemSize.height > 0) {
            return CTFloorCGFloat(itemSize.height * columnWidth / itemSize.width);
        }
    }

    return 0.0f;
}

- (CGFloat)m_heightForHeaderInSection:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:heightForHeaderInSection:)]) {
        return [self.delegate collectionView:self.collectionView layout:self heightForHeaderInSection:section];
    } else {
        return self.headerHeight;
    }
}

- (CGFloat)m_heightForFooterInSection:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:heightForFooterInSection:)]) {
        return [self.delegate collectionView:self.collectionView layout:self heightForFooterInSection:section];
    } else {
        return self.footerHeight;
    }
}

- (UIEdgeInsets)m_insetForSectionAtIndex:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:insetForSectionAtIndex:)]) {
        return [self.delegate collectionView:self.collectionView layout:self insetForSectionAtIndex:section];
    } else {
        return self.sectionInset;
    }
}

- (UIEdgeInsets)m_insetForHeaderInSection:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:insetForHeaderInSection:)]) {
        return [self.delegate collectionView:self.collectionView layout:self insetForHeaderInSection:section];
    } else {
        return self.headerInset;
    }
}

- (UIEdgeInsets)m_insetForFooterInSection:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:insetForFooterInSection:)]) {
        return [self.delegate collectionView:self.collectionView layout:self insetForFooterInSection:section];
    } else {
        return self.footerInset;
    }
}

- (CGFloat)m_interitemSpacingForSectionAtIndex:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:interitemSpacingForSectionAtIndex:)]) {
        return [self.delegate collectionView:self.collectionView layout:self interitemSpacingForSectionAtIndex:section];
    } else {
        return self.interitemSpacing;
    }
}

- (CGFloat)m_columnSpacingForSectionAtIndex:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:columnSpacingForSectionAtIndex:)]) {
        return [self.delegate collectionView:self.collectionView layout:self columnSpacingForSectionAtIndex:section];
    } else {
        return self.columnSpacing;
    }
}

- (CTCollectionViewItemRenderDirection)m_itemRenderDirectionForSectionAtIndex:(NSInteger)section {
    if ([self.delegate respondsToSelector:@selector(collectionView:layout:itemRenderDirectionForSectionAtIndex:)]) {
        return [self.delegate collectionView:self.collectionView layout:self itemRenderDirectionForSectionAtIndex:section];
    } else {
        return self.itemRenderDirection;
    }
}

- (BOOL)m_shouldPinToVisibleBoundsForSupplementaryViewOfKind:(NSString *)kind atSection:(NSInteger)section
{
    if ([kind isEqualToString:CTCollectionElementKindSectionHeader]) {
        if ([self.delegate respondsToSelector:@selector(collectionView:layout:shouldPinToVisibleBoundsForSupplementaryViewOfKind:atSection:)]) {
            return [self.delegate collectionView:self.collectionView layout:self shouldPinToVisibleBoundsForSupplementaryViewOfKind:kind atSection:section];
        }
        return self.sectionHeadersPinToVisibleBounds;
    } else if ([kind isEqualToString:CTCollectionElementKindSectionFooter]) {
        if ([self.delegate respondsToSelector:@selector(collectionView:layout:shouldPinToVisibleBoundsForSupplementaryViewOfKind:atSection:)]) {
            return [self.delegate collectionView:self.collectionView layout:self shouldPinToVisibleBoundsForSupplementaryViewOfKind:kind atSection:section];
        }
        return self.sectionFootersPinToVisibleBounds;
    }
    return NO;
}

- (NSArray<UICollectionViewLayoutAttributes *> *)m_layoutAttributesForSupplementaryViewInRect:(CGRect)rect
{
    if (!self.collectionView || !_haveFloatingHeaderOrFooter) return nil;

    NSMutableArray<UICollectionViewLayoutAttributes *> *attributeList = [NSMutableArray new];
    NSMutableArray<NSNumber *> *allFloatingSectionIndexs = [NSMutableArray new];

    for (NSInteger section = 0; section < self.sectionInfos.count; section++) {
        CTCollectionViewSectionInfo *sectionInfo =  self.sectionInfos[section];
        if ([sectionInfo hasFloatingHeaderOrFooter]) {
            [allFloatingSectionIndexs addObject:@(section)];
        } else {
            // normal header footer
            if (sectionInfo.headerAttribute != nil && sectionInfo.headerHeight > 0) {
                CGRect headerFrame = sectionInfo.headerAttribute.frame;
                headerFrame.origin.y = sectionInfo.headerBounds.top;
                if (CGRectIntersectsRect(headerFrame, rect)) {
                    [attributeList addObject:sectionInfo.headerAttribute];
                }
            }

            if (sectionInfo.footerAttribute != nil && sectionInfo.footerHeight > 0) {
                CGRect footerFrame = sectionInfo.footerAttribute.frame;
                footerFrame.origin.y = sectionInfo.footerBounds.top;
                if (CGRectIntersectsRect(footerFrame, rect)) {
                    [attributeList addObject:sectionInfo.footerAttribute];
                }
            }
        }
    }

    if (allFloatingSectionIndexs.count < 1) {
        return attributeList;
    }

    CGFloat offsetY = self.collectionView.contentOffset.y;
    CGFloat top = offsetY;

    NSNumber *lastFloatingItemTop = nil;
    BOOL flag = NO;
    BOOL skip = NO;
    for (NSInteger idx = allFloatingSectionIndexs.count - 1; idx >= 0; idx--) {
        NSInteger section = allFloatingSectionIndexs[idx].integerValue;
        CTCollectionViewSectionInfo *sectionInfo = self.sectionInfos[section];

        if (sectionInfo.headerBounds.top >= top || skip) {
            if (sectionInfo.headerAttribute != nil) {
                [self.class updateLayoutAttributesTop:sectionInfo.headerAttribute top:sectionInfo.headerBounds.top];
                sectionInfo.headerAttribute.zIndex = 0;
                [attributeList addObject:sectionInfo.headerAttribute];

                lastFloatingItemTop = @(CGRectGetMinY(sectionInfo.headerAttribute.frame));
            }
        } else {
            if (sectionInfo.headerAttribute != nil) {
                CGFloat top = MAX(0, offsetY);
                if (lastFloatingItemTop != nil && top + sectionInfo.headerHeight > lastFloatingItemTop.floatValue) {
                    top = lastFloatingItemTop.floatValue - sectionInfo.headerHeight;
                    flag = YES;

                    [self.class updateLayoutAttributesTop:sectionInfo.headerAttribute top:top];
                    sectionInfo.headerAttribute.zIndex = 1024 + section;
                    [attributeList addObject:sectionInfo.headerAttribute];
                } else {
                    if (flag) {
                        [self.class updateLayoutAttributesTop:sectionInfo.headerAttribute top:sectionInfo.headerBounds.top];
                        sectionInfo.headerAttribute.zIndex = 0;
                        [attributeList addObject:sectionInfo.headerAttribute];
                    } else {
                        [self.class updateLayoutAttributesTop:sectionInfo.headerAttribute top:top];
                        sectionInfo.headerAttribute.zIndex = 1024 + section;
                        [attributeList addObject:sectionInfo.headerAttribute];
                    }
                    skip = YES;
                }
            }
        }
    }

    return attributeList;
}

#pragma mark - Init

- (instancetype)init {
    if (self = [super init]) {
        [self m_configDefaultValue];
    }
    return self;
}

- (instancetype)initWithCoder:(NSCoder *)aDecoder {
    if (self = [super initWithCoder:aDecoder]) {
        [self m_configDefaultValue];
    }
    return self;
}

- (void)m_configDefaultValue {
    _columnCount = 2;
    _columnSpacing = 10;
    _interitemSpacing = 10;
    _headerHeight = 0;
    _footerHeight = 0;
    _sectionInset = UIEdgeInsetsZero;
    _headerInset = UIEdgeInsetsZero;
    _footerInset = UIEdgeInsetsZero;
    _itemRenderDirection = CTCollectionViewItemRenderDirectionShortestFirst;
    _allowUseSafeAreaInsets = YES;
    _sectionHeadersPinToVisibleBounds = NO;
    _sectionFootersPinToVisibleBounds = NO;
}

#pragma mark - Override Methods

- (void)prepareLayout {
    [super prepareLayout];

    [self.sectionInfos removeAllObjects];
    _haveFloatingHeaderOrFooter = NO;

    NSInteger numberOfSections = [self.collectionView numberOfSections];
    if (numberOfSections == 0) {
        return;
    }

    NSAssert([self.delegate conformsToProtocol:@protocol(CTCollectionViewDelegateColumnLayout)], @"UICollectionView's delegate should conform to CTCollectionViewDelegateColumnLayout protocol");
    NSAssert(self.columnCount > 0 || [self.delegate respondsToSelector:@selector(collectionView:layout:columnCountForSection:)], @"CTCollectionViewColumnLayout's columnCount should be greater than 0, or delegate must implement -columnCountForSection:");

    CGFloat top = 0.0;

    UIEdgeInsets safeAreaInsets = UIEdgeInsetsZero;
    if (self.allowUseSafeAreaInsets) {
        if (@available(iOS 11.0, *)) {
            safeAreaInsets = self.collectionView.safeAreaInsets;
        }
    }

    for (NSInteger section = 0; section < numberOfSections; section++) {
        CTCollectionViewSectionInfo *sectionInfo = [[CTCollectionViewSectionInfo alloc] init];
        sectionInfo.section = section;
        sectionInfo.columnCount = [self m_columnCountForSection:section];

        sectionInfo.columnSpacing = [self m_columnSpacingForSectionAtIndex:section];
        sectionInfo.interitemSpacing = [self m_interitemSpacingForSectionAtIndex:section];
        sectionInfo.collectionViewWidth = CGRectGetWidth(self.collectionView.bounds);
        sectionInfo.columnWidth = [self itemWidthForSection:section];

        sectionInfo.top = top;
        sectionInfo.headerHeight = [self m_heightForHeaderInSection:section];
        sectionInfo.footerHeight = [self m_heightForFooterInSection:section];
        sectionInfo.headerInsets = [self m_insetForHeaderInSection:section];
        sectionInfo.footerInsets = [self m_insetForFooterInSection:section];
        sectionInfo.sectionInsets = [self m_insetForSectionAtIndex:section];
        sectionInfo.safeAreaInsets = safeAreaInsets;
        sectionInfo.itemRenderDirection = [self m_itemRenderDirectionForSectionAtIndex:section];

        if (sectionInfo.headerHeight > 0) {
            sectionInfo.headerPinToVisible = [self m_shouldPinToVisibleBoundsForSupplementaryViewOfKind:CTCollectionElementKindSectionHeader atSection:section];
        } else {
            sectionInfo.headerPinToVisible = NO;
        }
        if (sectionInfo.footerHeight > 0) {
            sectionInfo.footerPinToVisible = [self m_shouldPinToVisibleBoundsForSupplementaryViewOfKind:CTCollectionElementKindSectionFooter atSection:section];
        } else {
            sectionInfo.footerPinToVisible = NO;
        }

        if (!_haveFloatingHeaderOrFooter && sectionInfo.hasFloatingHeaderOrFooter) {
            _haveFloatingHeaderOrFooter = YES;
        }

        NSInteger itemCount = [self.collectionView numberOfItemsInSection:section];
        for (NSInteger idx = 0; idx < itemCount; idx++) {
            CTCollectionViewItemInfo *itemInfo = [[CTCollectionViewItemInfo alloc] init];
            itemInfo.section = section;
            itemInfo.item = idx;
            itemInfo.columnIndex = 0;
            itemInfo.itemHeight = [self m_heightForItemAtIndexPath:[NSIndexPath indexPathForItem:idx inSection:section] columnWidth:sectionInfo.columnWidth];

            [sectionInfo.allItemInfos addObject:itemInfo];
        }

        // invoke this method will calculate all layout infos.
        top += [sectionInfo totalSectionHeight];

        [self.sectionInfos addObject:sectionInfo];
    }
}

- (CGSize)collectionViewContentSize {
    NSInteger numberOfSections = [self.collectionView numberOfSections];
    if (numberOfSections == 0) {
        return CGSizeZero;
    }

    CGSize contentSize = CGSizeMake(self.collectionView.bounds.size.width, 0.0);
    if (self.sectionInfos.count > 0) {
        contentSize.height = [self.sectionInfos lastObject].sectionBottom;
    }

    if (contentSize.height < self.minimumContentHeight) {
        contentSize.height = self.minimumContentHeight;
    }
    return contentSize;
}

- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath {
    if (indexPath.section >= 0 && indexPath.section < self.sectionInfos.count) {
        CTCollectionViewSectionInfo *sectionInfo = [self.sectionInfos objectAtIndex:indexPath.section];
        if (indexPath.item >= 0 && indexPath.item < sectionInfo.allItemInfos.count) {
            return sectionInfo.allItemInfos[indexPath.item].itemAttributes;
        }
    }
    return nil;
}

- (UICollectionViewLayoutAttributes *)layoutAttributesForSupplementaryViewOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath {
    if (indexPath.section >= 0 && indexPath.section < self.sectionInfos.count) {
        CTCollectionViewSectionInfo *sectionInfo = self.sectionInfos[indexPath.section];

        if ([kind isEqualToString:CTCollectionElementKindSectionHeader]) {
            [self.class updateLayoutAttributesTop:sectionInfo.headerAttribute top:sectionInfo.headerBounds.top];
            return sectionInfo.headerAttribute;
        } else if ([kind isEqualToString:CTCollectionElementKindSectionFooter]) {
            [self.class updateLayoutAttributesTop:sectionInfo.footerAttribute top:sectionInfo.footerBounds.top];
            return sectionInfo.footerAttribute;
        }
    }
    return nil;
}

- (NSArray<__kindof UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect {
    NSMutableArray *attributesList = [[NSMutableArray alloc] initWithCapacity:0];

    BOOL stop = NO;
    BOOL exists = NO;
    for (CTCollectionViewSectionInfo *sectionInfo in self.sectionInfos) {
        if (stop) {
            break;
        }

        BOOL hasHeader = sectionInfo.headerHeight > 0 ? YES : NO;
        BOOL hasFooter = sectionInfo.footerHeight > 0 ? YES : NO;

        if (!exists) {
            if (sectionInfo.footerAttribute != nil && hasFooter && CGRectGetMinY(rect) > CGRectGetMaxY(sectionInfo.footerAttribute.frame)) {
                continue;
            }
        } else {
            if (sectionInfo.headerAttribute != nil && hasHeader && CGRectGetMaxY(rect) < CGRectGetMinY(sectionInfo.headerAttribute.frame)) {
                stop = YES;
                break;
            }
        }

        if (sectionInfo.headerAttribute != nil) {
            CGRect headerFrame = sectionInfo.headerAttribute.frame;
            headerFrame.origin.y = sectionInfo.headerBounds.top;

            if (hasHeader && CGRectIntersectsRect(headerFrame, rect)) {
                exists = YES;
                if (!_haveFloatingHeaderOrFooter) {
                    sectionInfo.headerAttribute.frame = headerFrame;
                    [attributesList addObject:sectionInfo.headerAttribute];
                }
            }
        }

        if (sectionInfo.footerAttribute != nil) {
            CGRect footerFrame = sectionInfo.footerAttribute.frame;
            footerFrame.origin.y = sectionInfo.footerBounds.top;

            if (hasFooter && CGRectIntersectsRect(footerFrame, rect)) {
                exists = YES;
                if (!_haveFloatingHeaderOrFooter) {
                    sectionInfo.footerAttribute.frame = footerFrame;
                    [attributesList addObject:sectionInfo.footerAttribute];
                }
            }
        }

        for (CTCollectionViewItemInfo *itemInfo in sectionInfo.allItemInfos) {
            if (itemInfo.itemAttributes != nil && CGRectIntersectsRect(itemInfo.itemAttributes.frame, rect)) {
                exists = YES;
                [attributesList addObject:itemInfo.itemAttributes];
            }
        }
    }

    if (_haveFloatingHeaderOrFooter) {
        NSArray *headerFooterAttributesList = [self m_layoutAttributesForSupplementaryViewInRect:rect];
        if (headerFooterAttributesList && headerFooterAttributesList.count > 0) {
            [attributesList addObjectsFromArray:headerFooterAttributesList];
        }
    }

    return attributesList;
}

- (BOOL)shouldInvalidateLayoutForBoundsChange:(CGRect)newBounds {
    if (!self.collectionView) {
        return [super shouldInvalidateLayoutForBoundsChange:newBounds];
    }

    if (_haveFloatingHeaderOrFooter) {
        return YES;
    }

    CGRect oldBounds = self.collectionView.bounds;
    if (CGRectGetWidth(newBounds) != CGRectGetWidth(oldBounds)) {
        return YES;
    }
    return NO;
}

+ (Class)layoutAttributesClass {
    return [CTCollectionViewColumnLayout class];
}

#pragma mark - Util

// update top of the frame for layout attributes
+ (void)updateLayoutAttributesTop:(UICollectionViewLayoutAttributes *)layoutAttributes top:(CGFloat)top
{
    if (!layoutAttributes) return;

    if (CGRectGetMinY(layoutAttributes.frame) != top) {
        CGRect frame = layoutAttributes.frame;
        frame.origin.y = top;
        layoutAttributes.frame = frame;
    }
}

@end
